package cas.ihep.fs.alluxio;

import alluxio.AlluxioURI;
import alluxio.Configuration;
import alluxio.PropertyKey;
import alluxio.client.block.BlockMasterClient;
import alluxio.client.block.policy.BlockLocationPolicy;
import alluxio.client.block.policy.options.GetWorkerOptions;
import alluxio.client.block.stream.BlockInStream;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.FileSystemMasterClient;
import alluxio.client.file.URIStatus;
import alluxio.client.file.options.*;
import alluxio.client.file.policy.FileWriteLocationPolicy;
import alluxio.exception.AlluxioException;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.PreconditionMessage;
import alluxio.exception.status.NotFoundException;
import alluxio.exception.status.UnavailableException;
import alluxio.master.block.BlockId;
import alluxio.proto.dataserver.Protocol;
import alluxio.resource.CloseableResource;
import alluxio.security.authorization.Mode;
import alluxio.util.CommonUtils;
import alluxio.util.network.NettyUtils;
import alluxio.wire.BlockInfo;
import alluxio.wire.BlockLocation;
import alluxio.wire.LoadMetadataType;
import alluxio.wire.WorkerNetAddress;
import cas.ihep.alluxio.AlluxioSystem;
import cas.ihep.fs.UnifiedBlock;
import cas.ihep.fs.UnifiedFile;
import cas.ihep.fs.UnifiedFileSystem;
import cas.ihep.util.AlluxioRPCContext;
import cas.ihep.util.Finalizer;
import cas.ihep.util.MemoryBuffer;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.InvocationTargetException;
import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

abstract class AlluxioFile implements UnifiedFile,Closeable{
    AlluxioFileSystem alluxioFileSystem;
    AlluxioURI alluxioUri;
    URIStatus uriStatus;
    List<AlluxioBlock> subBlocks;
    Finalizer closer;
    long fileLength;
    long blockSize;
    boolean isClosed;

    AlluxioFile(AlluxioFileSystem sys,AlluxioURI uri){
        alluxioFileSystem=sys;
        alluxioUri=uri;
        uriStatus=null;
        closer=Finalizer.create();
        fileLength=0;
        blockSize=0;
        isClosed=false;
    }

    public long length(){
        return fileLength;
    }

    public URI path(){
        return AlluxioFileSystem.baseURI(alluxioUri);
    }

    long getBlockSize(long pos){
        long lastBlockSize = fileLength % blockSize;
        if (fileLength - pos > lastBlockSize) {
            return blockSize;
        } else {
            return lastBlockSize;
        }
    }

    public void free()throws IOException{
        try {
            close();
            alluxioFileSystem.fileSystem.free(alluxioUri);
        } catch (AlluxioException e) {
            throw new IOException(e);
        }
    }

    public void persist()throws IOException{
        if(uriStatus.isPersisted()){
            return;
        }
        close();
        try(CloseableResource<FileSystemMasterClient> client=alluxioFileSystem.fileSystemContext.acquireMasterClientResource()){
            client.get().scheduleAsyncPersist(alluxioUri);
        }
        URIStatus newStatus;
        do{
            try {
                newStatus=alluxioFileSystem.fileSystem.getStatus(alluxioUri);
                if(newStatus.isPersisted()){
                    break;
                }
                Thread.sleep(1000);
            } catch (InterruptedException|AlluxioException ignored) {
                return;
            }
        }while(true);
        uriStatus=newStatus;
        //FileSystemUtils.persistFile(alluxioFileSystem.fileSystem,alluxioUri);
    }

    public void delete()throws IOException{
        close();
        try {
            alluxioFileSystem.fileSystem.delete(alluxioUri);
        } catch (AlluxioException e) {
            throw new IOException(e);
        }
    }

    static class AlluxioFileIn extends AlluxioFile{
        InStreamOptions mInOptions;
        private OutStreamOptions mOutOptions;

        AlluxioFileIn(AlluxioFileSystem sys, AlluxioURI uri) throws IOException{
            super(sys, uri);
            mInOptions=InStreamOptions.defaults();
            mOutOptions=OutStreamOptions.defaults();
            try{
                URIStatus status=super.alluxioFileSystem.fileSystem.getStatus(super.alluxioUri);
                super.blockSize=status.getBlockSizeBytes();
                super.fileLength=status.getLength();
                super.uriStatus=status;
                int sz=status.getBlockIds().size();
                super.subBlocks=new ArrayList<>(sz);
                for(int i=0;i<sz;i++){
                    subBlocks.add(null);
                }
            }catch (AlluxioException e){
                throw new IOException(e);
            }
        }

        public void close()throws IOException{
            if(!isClosed){
                isClosed=true;
                closer.registerAll(subBlocks,false);
                closer.close();
            }
        }

        @Override
        public int permission() {
            return PERMISSION_READ;
        }

        @Override
        public UnifiedBlock block(int index) throws IOException {
            AlluxioBlock blk=subBlocks.get(index);
            if(blk==null){
                long id=uriStatus.getBlockIds().get(index);
                blk = findBlock(id);
                subBlocks.set(index,blk);
            }
            return blk;
        }

        @Override
        public UnifiedBlock next() throws IOException {
            throw new IOException("Unsupported operation");
        }

        private AlluxioBlock findBlock(long blockId)throws IOException {
            AlluxioInvoker handler=alluxioFileSystem.alluxioInvoker;
            FileSystemContext context=alluxioFileSystem.fileSystemContext;
            Protocol.OpenUfsBlockOptions openUfsBlockOptions=null;
            long mBlockSize=super.blockSize;
            URIStatus mStatus=super.uriStatus;
            if (mStatus.isPersisted()) {
                long blockStart = BlockId.getSequenceNumber(blockId) * mBlockSize;
                openUfsBlockOptions =
                        Protocol.OpenUfsBlockOptions.newBuilder().setUfsPath(mStatus.getUfsPath())
                                .setOffsetInFile(blockStart).setBlockSize(getBlockSize(blockStart))
                                .setMaxUfsReadConcurrency(mInOptions.getMaxUfsReadConcurrency())
                                .setNoCache(!mInOptions.getAlluxioStorageType().isStore())
                                .setMountId(mStatus.getMountId()).build();
            }
            BlockInfo blockInfo;
            try (CloseableResource<BlockMasterClient> masterClientResource =
                         context.acquireBlockMasterClientResource()) {
                blockInfo = masterClientResource.get().getBlockInfo(blockId);
            }
            List<BlockLocation> locations=blockInfo.getLocations();
            if (locations.isEmpty() && openUfsBlockOptions == null) {
                throw new NotFoundException("Block " + blockId + " is unavailable in both Alluxio and UFS.");
            }
            long blockSize=blockInfo.getLength();
            for (BlockLocation location : locations) {
                WorkerNetAddress workerNetAddress = location.getWorkerAddress();
                if (workerNetAddress.getHost().equals(AlluxioSystem.localHostName)) {
                    return new AlluxioBlock.AlluxioReadableBlock(this,workerNetAddress,blockId,blockSize);
                }
            }
            BlockInStream.BlockInStreamSource source = BlockInStream.BlockInStreamSource.UFS;
            WorkerNetAddress address;
            if (locations.isEmpty()) {
                BlockLocationPolicy blockLocationPolicy =
                //        Preconditions.checkNotNull(mInOptions.getUfsReadLocationPolicy(),
                //                PreconditionMessage.UFS_READ_LOCATION_POLICY_UNSPECIFIED);
                    Objects.requireNonNull(mInOptions.getUfsReadLocationPolicy(),PreconditionMessage.UFS_READ_LOCATION_POLICY_UNSPECIFIED.toString());
                address = blockLocationPolicy
                        .getWorker(GetWorkerOptions.defaults().setBlockWorkerInfos(handler.getWorkerInfoList())
                                .setBlockId(blockId).setBlockSize(blockInfo.getLength()));
                if (address == null) {
                    throw new UnavailableException(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage());
                }
            } else {
                address = locations.get(handler.nextInt(locations.size())).getWorkerAddress();
                source = BlockInStream.BlockInStreamSource.REMOTE;
            }
            Protocol.ReadRequest.Builder builder = Protocol.ReadRequest.newBuilder().setBlockId(blockId)
                    .setPromote(mInOptions.getAlluxioStorageType().isPromote());
            if (openUfsBlockOptions != null) {
                builder.setOpenUfsBlockOptions(openUfsBlockOptions);
            }
            WorkerNetAddress localWorker=context.getLocalWorker();
            if (localWorker!=null) {
                AlluxioRPCContext rpcContext;
                try (Finalizer closer = Finalizer.create()) {
                    InputStream is = closer.register(handler.createNettyBlockInStream(context, address, source, builder.buildPartial(), blockSize, mInOptions));
                    rpcContext=closer.register(new AlluxioRPCContext(alluxioFileSystem.fileSystemContext,localWorker));
                    AlluxioBlock.AlluxioWritableBlock tmpBlock = closer.register(new AlluxioBlock.AlluxioWritableBlock(rpcContext, blockId, blockSize, mOutOptions));
                    tmpBlock.needCloseOrCancel = false;
                    MemoryBuffer buffer = closer.register(new MemoryBuffer(new File(tmpBlock.blockPath), 0, blockSize));
                    byte bbuf[] = new byte[8192];
                    for (long total = 0; total < blockSize; ) {
                        int each = is.read(bbuf);
                        buffer.put(bbuf, 0, each);
                        total += each;
                    }
                    tmpBlock.needCloseOrCancel = true;
                    closer.cancel(rpcContext);
                } catch (IllegalAccessException e) {
                    throw new IOException(e.getMessage());
                } catch (InvocationTargetException e) {
                    throw new IOException(e.getTargetException());
                }
                return new AlluxioBlock.AlluxioReadableBlock(rpcContext, blockId, blockSize,mInOptions).manageRPC();
            }
            throw new IOException("Cannot find block on local worker");
//            try (InputStream is = handler.createNettyBlockInStream(context, address, source,
//                    builder.buildPartial(), blockSize, mInOptions)){
//                WorkerNetAddress localWorker = context.getLocalWorker();
//                if (localWorker != null) {
//                    try (OutputStream os =
//                                 handler.getOutStream(blockId, blockSize, localWorker, mOutOptions)) {
//                        ByteStreams.copy(is,os);
//                    }
//                    return new AlluxioBlock.AlluxioReadableBlock(this,localWorker,blockId,blockSize);
//
//                }
//                throw new IOException("Cannot find block on local worker");
//            }catch (IllegalAccessException | InvocationTargetException e){
//                throw new IOException(e);
//            }
        }

    }

    static class AlluxioFileOut extends AlluxioFile{
        private OutStreamOptions mCreateOptions;
        private FileSystemMasterClient masterClient;

        public void close()throws IOException{
            if(!isClosed){
                isClosed=true;
                closer.register(new Closeable() {
                    @Override
                    public void close() throws IOException {
                        URIStatus status;
                        GetStatusOptions options=GetStatusOptions.defaults().setLoadMetadataType(LoadMetadataType.Never);
                        int i=0;
                        do{
                            if(i++>10){
                                throw new IOException("blocks mismatch");
                            }
                            status=masterClient.getStatus(alluxioUri,options);
                            if(status.getBlockIds().size()==subBlocks.size()){
                                break;
                            }
                            try {
                                Thread.sleep(100);
                            } catch (InterruptedException ignored) {}
                        }while(true);
                        masterClient.completeFile(alluxioUri, CompleteFileOptions.defaults());
                    }
                });
                closer.registerAll(subBlocks,true);
                closer.close();
            }
        }

        AlluxioFileOut(AlluxioFileSystem sys, AlluxioURI uri, UnifiedFileSystem.CreateOptions copt) throws IOException {
            super(sys, uri);
            subBlocks=new ArrayList<>();
            CreateFileOptions options=CreateFileOptions.defaults();
            options.setMode(new Mode(copt.permission()));
            options.setRecursive(copt.createParent());
            options.setBlockSizeBytes(copt.blockSize());
            try {
                if (copt.overwrite() && sys.fileSystem.exists(uri)) {
                    sys.fileSystem.delete(uri);
                }
            }catch (AlluxioException e){
                throw new IOException(e);
            }
            masterClient=alluxioFileSystem.fileSystemContext.acquireMasterClient();
            closer.register(new Closeable() {
                @Override
                public void close() {
                    alluxioFileSystem.fileSystemContext.releaseMasterClient(masterClient);
                }
            });
            masterClient.createFile(uri,options);
            URIStatus mStatus=masterClient.getStatus(uri,GetStatusOptions.defaults().setLoadMetadataType(LoadMetadataType.Never));
            mCreateOptions=options.toOutStreamOptions();
            mCreateOptions.setUfsPath(mStatus.getUfsPath());
            mCreateOptions.setMountId(mStatus.getMountId());
            super.blockSize=options.getBlockSizeBytes();
            super.uriStatus=mStatus;
        }

        @Override
        public int permission() {
            return PERMISSION_WRITE;
        }

        @Override
        public UnifiedBlock block(int index) {
            return subBlocks.get(index);
        }

        private AlluxioBlock nextBlock()throws IOException{
            long nextid=getNextBlockId();
            long mBlockSize=blockSize;
            AlluxioFileSystem fs=alluxioFileSystem;
            AlluxioInvoker invoker=fs.alluxioInvoker;
            FileSystemContext ctx=fs.fileSystemContext;
            WorkerNetAddress address;
            FileWriteLocationPolicy locationPolicy =
                    //Preconditions.checkNotNull(mCreateOptions.getLocationPolicy(),PreconditionMessage.FILE_WRITE_LOCATION_POLICY_UNSPECIFIED);
                Objects.requireNonNull(mCreateOptions.getLocationPolicy(),PreconditionMessage.FILE_WRITE_LOCATION_POLICY_UNSPECIFIED.toString());
            address = locationPolicy.getWorkerForNextBlock(invoker.getWorkerInfoList(), mBlockSize);
            if (address == null) {
                throw new UnavailableException(
                        ExceptionMessage.NO_SPACE_FOR_BLOCK_ON_WORKER.getMessage(mBlockSize));
            }
            if (mBlockSize == -1) {
                try (CloseableResource<BlockMasterClient> blockMasterClientResource =
                             ctx.acquireBlockMasterClientResource()) {
                    mBlockSize = blockMasterClientResource.get().getBlockInfo(nextid).getLength();
                }
            }
            if (CommonUtils.isLocalHost(address) && Configuration
                    .getBoolean(PropertyKey.USER_SHORT_CIRCUIT_ENABLED) && !NettyUtils
                    .isDomainSocketSupported(address)) {
                return new AlluxioBlock.AlluxioWritableBlock(this,address,nextid,mBlockSize,mCreateOptions);
            }
            throw new IOException("Unable to allocate local block");
        }

        @Override
        public UnifiedBlock next() throws IOException {
            AlluxioBlock blk=nextBlock();
            subBlocks.add(blk);
            return blk;
        }

        long getNextBlockId()throws IOException{
            return masterClient.getNewBlockIdForFile(alluxioUri);
        }
    }


}
